import sys
import os
from typing import Any, Dict, List, Optional
from atomic_agents.lib.base.base_tool import BaseTool, BaseToolConfig
from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator
from pydantic import BaseModel, ConfigDict, Field
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from BasicTools.BehaviorQuantModel.ChatRecordProcessor import ChatRecordProcessor
from BasicTools.BehaviorQuantModel.DataModels import ChatRecord,OprationChain
class GetOprationChainFromChatInput(BaseModel):
    """
    Input schema for GetOprationChainFromChat tool.
    """
    model_config=ConfigDict(
        json_schema_extra={
            "example": {
                "opration_user_id": "user_001",
                "opration_question_id": "q001",
                "action_type":"defense"
            }
        }
    )
    opration_user_id: str = Field(..., description="the user id of the opration agent",examples="user_001")
    opration_question_id: str = Field(..., description="the question id of the opration agent", examples="q001")
    action_type: str = Field(..., description="the action type of the opration agent", examples="attack")
class GetOprationChainFromChatOutput(BaseModel):
    """
    Output schema for GetOprationChainFromChat tool.
    """
    model_config=ConfigDict(
        json_schema_extra={
            "example": {
                "opration_chain": {
                    "user_id": "user_001",
                    "question_id": "q001",
                    "actions": ["action1", "action2", "action3"]        
                },
            }
        }
    )
    opration_chain: OprationChain = Field(..., description="the opration chain data extracted from chat record.",
            examples={
                "user_id":"user_001",
                "question_id":"q001",
                "actions":["action1", "action2", "action3"],
            })
    

class GetOprationChainTool(BaseTool):
    """
    Tool to get operation chain from chat record.
    Returns structured operation chain data and analysis results.
    """
    name = "GetOprationChain"
    description = "A tool to get operation chain from chat record."
    input_schema = GetOprationChainFromChatInput
    output_schema = GetOprationChainFromChatOutput
    
    def __init__(self, config: BaseToolConfig = None):
        super().__init__(config)
    
    def run(self, inputs: GetOprationChainFromChatInput) -> GetOprationChainFromChatOutput:
        """
        从聊天记录中提取攻击链
        Args:
            inputs (GetOprationChainFromChatInput): 输入参数
        Returns:
            GetOprationChainFromChatOutput: 攻击链数据和分析结果
        """
        chat_record = ChatRecordProcessor().get_formatted_chat_record(
            user_id=inputs.opration_user_id,
            question_id=inputs.opration_question_id,
            agent_type=inputs.action_type
        )
        opration_chain = chat_record.opration_chain
        return GetOprationChainFromChatOutput(opration_chain=opration_chain)
        